*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
*	This program estimates the RF and FS matrix.
*	----------------------------------------------------------------------------

	args sample sch_res ptype bw
	tokenize `sample', parse("_")
	local years "`5'`6'`7'`8'"

	// pick outcome test
	local y m

*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

*	load analysis file

	use "${builddata}${city}_analysisfile`sample'`sch_res'`bw'", clear
	merge m:1 sch year using "${builddata}${city}_schoolcovs", nogen keep(1 3)
	if inlist("${city}","NYC","NYCms") {
		ren screened screened_old
		qui mdesc screened
		assert r(miss) == 0
		egen any_scr = max(screened), by(sch)
		replace screened = any_scr
	}

	preserve

	ivset, ptype(`ptype') y(`y')
		local Dfull = r(D)
		local S = r(S)
		local p = r(P)
		local Y = r(Y)

	//need schools to be represented in the lottery sample, or else first step regs have no variance ests.
	_rmcoll `Dfull' if `S', forcedrop nocons
	local D = r(varlist)

	gen tag = 0
	foreach var of varlist `D'{
		local sch = substr("`var'",-4,.)
		replace tag = 1 if sch == `sch'
	}

	keep if !omitted & tag

	keep sch lottery_`ptype'_sch $sectors
	duplicates drop

	//lotteries first
	gsort -lottery_`ptype'_sch sch
	gen newnum = _n

	//count schools
	local J=_N
	count if lottery_`ptype'_sch
	local L=r(N)

	di in red "schools: `J'"
	di in red "lotteries: `L'"

	//store ordering in locals
	forval i=1/`J'{
		qui sum sch if newnum == `i'
		local oldnum_`i' = r(mean)
	}

***	stack key

	order lottery_`ptype'_sch $sectors
	save "${builddata}${city}_sortkey_`ptype'`sample'`sch_res'", replace

	drop sch newnum

	export delimited using "${builddata}${city}_sectors_`ptype'`sample'`sch_res'`bw'",  novarnames replace
	save "${builddata}${city}_stackkey_`ptype'`sample'`sch_res'`bw'", replace

	restore

*	OLS-VAM

	//enrollment
	forval j = 1/`J'{
		gen vam_`j' = enr_`oldnum_`j''
	}

	//het
	mat b1 = J(1,`J',.)
	mat V1homo = J(`J',`J',.)
	forval j = 1/`J'{
		di in red "OLS: `j'"
		qui _regress math vam_`j' $iv_controls,
		mat b1[1,`j'] = _b[vam_`j']
		mat V1homo[`j',`j'] = _se[vam_`j']^2
	}

*	reduced form

	//get corresponding programs for each school
	if strpos("`ptype'","2") {
		preserve
			clear
			set obs 1
			if "${city}"=="NYC" {
				forval year = 2012/2014{
					append using "${cleandata}/prg_level_pscores`bw'_`year'.dta"
				}
			}
			if "${city}"=="NYCms" {
				forval year = 2016/2018{
					append using "${cleandata}/prg_level_pscores`bw'_ms_`year'.dta"
				}
				ren programcode prg
			}
			drop in 1
			gen sch = substr(sch_long, 3, 4)
			//destring school codes
			replace sch = subinstr(sch,"B","1",.)
			replace sch = subinstr(sch,"K","2",.)
			replace sch = subinstr(sch,"M","3",.)
			replace sch = subinstr(sch,"Q","4",.)
			replace sch = subinstr(sch,"R","5",.)
			replace sch = subinstr(sch,"X","6",.)
			destring sch, replace

			levelsof sch, local(schools)
			foreach sch in `schools'{
				levelsof prg if sch == `sch', local(prgs_`sch')
			}
		restore
	}

	//offers
	forval l = 1/`L'{
		di in red "RF prep: `l'"
		gen rf_`l' = offer_`oldnum_`l''
		gen pscore_`l' = pscore_`ptype'_`oldnum_`l''
		gen pscore_`l'_0 = (pscore_`l' == 0)
		if strpos("`ptype'","2") | strpos("`ptype'","3"){
			foreach prg in `prgs_`oldnum_`l'''{
				local rvs
				cap unab rvs: rv_*`prg'*
				local rv_`l' `rv_`l'' `rvs'
			}
		}
		di "rv `l': `rv_`l''"
	}

	//rf
	preserve
		ivset,
			local Dfull enr_????
			local S = r(S)

			local rf
			local p
			local rv
			forval l = 1/`L'{
				local rf `rf' rf_`l'
				local p `p' pscore_`l' pscore_`l'_0
				local rv `rv' `rv_`l''
				di `rv_`l''
			}
		qui _regress `Y' `rf' `p' $iv_controls i.year#c.(`rv') if `S'
		mat V = e(V)
		mat V = V[1..`L',1..`L']
		clear
		svmat V
		if "`y'" == "m" export delimited using "${builddata}${city}_rf_V_1.csv", novarnames replace
		else export delimited using "${builddata}${city}_rf_V_0.csv", novarnames replace
	restore

	mat b2 = J(1,`L',.)
	mat V2homo = J(`L',`L',.)
	forval l = 1/`L'{
		di in red "RF: `l'"
 		qui _regress `Y' rf_`l' pscore_`l' pscore_`l'_0 i.year#c.(`rv_`l'') $iv_controls if `S',
		mat b2[1,`l'] = _b[rf_`l']
		mat V2homo[`l',`l'] = _se[rf_`l']^2
	}

	if "`y'" == "e" {
		forval m = 1/2{
			preserve
			clear
			svmat2 b`m'
			export delimited "${builddata}diagdiag_e/${city}_`ptype'`bw'_b`m'.csv", replace novarnames
			clear
			svmat2 V`m'homo
			export delimited "${builddata}diagdiag_e/${city}_`ptype'`bw'_V`m'homo.csv", replace novarnames
			restore
		}
	}
	else if "`y'" == "m" {
		forval m = 1/2{
			preserve
			clear
			svmat2 b`m'
			export delimited "${builddata}diagdiag_m/${city}_`ptype'`bw'_b`m'.csv", replace novarnames
			clear
			svmat2 V`m'homo
			export delimited "${builddata}diagdiag_m/${city}_`ptype'`bw'_V`m'homo.csv", replace novarnames
			restore
		}
	}

*	first stages

	//first stages
	local i 2
	forval j = 1/`J'{
		local i = `i' + 1
		//homo
		mat b`i' = J(1,`L',.)
		mat V`i'homo = J(`L',`L',.)
		forval l = 1/`L'{
			di in red "FS: `j', `l'"
			qui _regress vam_`j' rf_`l' pscore_`l'  pscore_`l'_0 `rv_`l'' $iv_controls if `S',
			mat b`i'[1,`l'] = _b[rf_`l']
			mat V`i'homo[`l',`l'] = _se[rf_`l']^2
		}

		if "`y'" == "e" {
			forval m = `=2+`j''/`=2+`j''{
				preserve
				clear
				svmat2 b`m'
				export delimited "${builddata}diagdiag_e/${city}_`ptype'`bw'_b`m'.csv", replace novarnames
				clear
				svmat2 V`m'homo
				export delimited "${builddata}diagdiag_e/${city}_`ptype'`bw'_V`m'homo.csv", replace novarnames
				restore
			}
		}
		else if "`y'" == "m" {
			forval m = `=2+`j''/`=2+`j''{
				preserve
				clear
				svmat2 b`m'
				export delimited "${builddata}diagdiag_m/${city}_`ptype'`bw'_b`m'.csv", replace novarnames
				clear
				svmat2 V`m'homo
				export delimited "${builddata}diagdiag_m/${city}_`ptype'`bw'_V`m'homo.csv", replace novarnames
				restore
			}
		}

	}
